package com.btl.service.ai.support.rag.embedding.strategy.retriever;

import com.btl.service.ai.support.rag.embedding.strategy.EmbeddingModelTypeEnum;
import com.btl.service.ai.support.rag.embedding.strategy.RetrieverStrategy;
import com.btl.service.ai.support.rag.store.vector.collections.VectorStoreQdrantFactory;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.filter.Filter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;

import java.util.List;

@Slf4j
@Service
public class BgeSmallZhV15RetrieverStrategy implements RetrieverStrategy {

    private final EmbeddingModel begSmallZhV15Model;
    private final VectorStoreQdrantFactory vectorStoreQdrantFactory;

    public BgeSmallZhV15RetrieverStrategy(@Qualifier("begSmallZhV15Model") EmbeddingModel begSmallZhV15Model, VectorStoreQdrantFactory vectorStoreQdrantFactory) {
        this.begSmallZhV15Model = begSmallZhV15Model;
        this.vectorStoreQdrantFactory = vectorStoreQdrantFactory;
    }

    @Override
    public List<Content> retriever(String collectionName, Double minScore, Integer maxResults, Filter filter, Query query) {
        EmbeddingStore<TextSegment> embeddingStoreQdrant = vectorStoreQdrantFactory.getQdrantClient(collectionName.concat("_")
                .concat(String.valueOf(begSmallZhV15Model.dimension())));
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStoreQdrant)
                .embeddingModel(begSmallZhV15Model)
                .maxResults(maxResults)
                .minScore(minScore)
                .filter(filter)
                .build();
        return contentRetriever.retrieve(query);
    }

    @Override
    public List<Content> retriever(String collectionName, Query query) {
        return retriever(collectionName, 0.75, 5, null, query);
    }

    @Override
    public List<Content> retriever(String collectionName, Filter filter, Query query) {
        return retriever(collectionName, 0.75, 5, filter, query);
    }


    @Override
    public EmbeddingModelTypeEnum modelType() {
        return EmbeddingModelTypeEnum.LOCAL_BGE_SMALL_ZH_V15_EMBEDDING_MODEL;
    }
}
